# file:test.py
import random
import torch
import pylab
import matplotlib.pyplot as plt
from Dataloader import get_train_data
from CNN import ConvNet

# 实例化模型
test_net = ConvNet()
# 读取已保存的模型
test_net.load_state_dict(torch.load('./cnnmodel.pth'))

# 随机生产一张图片
x = random.randint(0, 59999)
train_data, _ = get_train_data()

# 丢入网络预测
output = test_net(train_data[x][0].unsqueeze(0))
pred = torch.argmax(output)
print(f"这张图片被识别为数字{pred}.")
print(f"这张图片实际为{train_data[x][1]}.")
if pred == train_data[x][1]:
    print("识别正确!")
else:
    print("识别错误!")

# 显示图像
plt.imshow(train_data[x][0].squeeze(0))
pylab.show()
